import os
import numpy as np
import pandas as pd
from FSdata.FSdataset import FSdata, collate_fn, idx2attr_map, attr2length_map
import torch
import torch.utils.data as torchdata
from models.resnet import resnet50_cat,resnet101_cat,resnet101_cat_dilate
from FSdata.FSaug import  *
from utils.preprocessing import addFakeLabel
from utils.predicting import predict,gen_submission
import zipfile
from utils.preprocessing import *

class FSAugTest(object):
    def __init__(self,doHflip):
        self.augment = Compose([
            ExpandBorder(select=[0, 5, 6, 7], mode='constant',resize=True,size=(336,336)),
            UpperCrop(size=(336, 336), select=[1,2,3,4]),
            Resize(size=(336, 336), select=[0, 1, 2, 3, 4, 5, 6, 7]),
            Hflip(doHflip=doHflip),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def __call__(self, image,attr_idx):
        return self.augment(image,attr_idx)


os.environ["CUDA_VISIBLE_DEVICES"] = "2"

rawdata_root = '/media/gserver/data/FashionAI'

test_root = os.path.join(rawdata_root, 'round2/week-rank')
test_pd = pd.read_csv(os.path.join(rawdata_root, 'round2/week-rank/Tests/question.csv'),
                       header=None, names=['ImageName', 'AttrKey', 'AttrValues'])


test_pd = addFakeLabel(test_pd)

# select part
select_AttrIdx = [1,2,3,4]
select_AttrKey = [idx2attr_map[x] for x in select_AttrIdx]
test_pd = test_pd[test_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]
rel_paths = test_pd['ImageName'].tolist()

test_pd = join_path_to_df(test_pd, rawdata_root, 'round2/week-rank')


print test_pd.shape
print test_pd['AttrKey'].value_counts()


# model prepare
model_name = 'res101-[1234]-9623-aug2-9677(dilate-pred)'
resume = '/media/gserver/extra/FashionAI/round2/res101_[1234]_base/bestweights-13-4458-[0.8670]-[0.9623].pth'
model = resnet101_cat_dilate(pretrained=True, num_classes=[attr2length_map[x] for x in select_AttrIdx])

# model = torch.nn.DataParallel(model)
print('resuming finetune from %s'%resume)
model.load_state_dict(torch.load(resume))
model = model.cuda()

# predict
if not os.path.exists('./online_pred2/csv'):
    os.makedirs('./online_pred2/csv')
if not os.path.exists('./online_pred2/subs_zip'):
    os.makedirs('./online_pred2/subs_zip')
if not os.path.exists('./online_pred2/part_csv'):
    os.makedirs('./online_pred2/part_csv')

aug_scores = np.zeros((2,test_pd.shape[0], sum([attr2length_map[x] for x in select_AttrIdx])),dtype=np.float32)

for i,doHflip in enumerate([0,1]):
    data_set = {}
    data_set['test'] = FSdata(anno_pd=test_pd,
                               transforms=FSAugTest(doHflip=doHflip),
                              select=select_AttrIdx
                               )

    data_loader = {}
    data_loader['test'] = torchdata.DataLoader(data_set['test'], batch_size=4, num_workers=4,
                                              shuffle=False, pin_memory=True,collate_fn=collate_fn)

    aug_scores[i,:,:],_,_, _ = predict(model, data_set['test'], data_loader['test'])


pred_scores = aug_scores.mean(0)

# save csv and scores
test_pred = gen_submission(test_pd[['ImageName', 'AttrKey']], pred_scores, catidx_map=data_set['test'].catidx_map)
test_pred['ImageName'] = rel_paths

if len(select_AttrIdx) < 8:
    test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred2/part_csv/%s.csv' % model_name,
                                                                 header=None, index=False)
else:
    test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred2/csv/%s.csv' % model_name,
                                                                 header=None, index=False)
    # make zip file
    z = zipfile.ZipFile('./online_pred2/subs_zip/%s.zip'%model_name, 'w', zipfile.ZIP_DEFLATED)
    z.write('./online_pred2/csv/%s.csv' % model_name, arcname='%s.csv' % model_name)
    z.close()


print test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].info()